from PIL import Image
import numpy as np
import os

FILE_NAME = os.path.join('BITDataeSet_images', 'multiple', '1.png')
IMG_SIZE = (32, 32)


class ImageError(Exception):
    __metaclass__ = object

    def __init__(self, *args, **kw):
        text = 'The image seems not like a digit.'
        super(ImageError, self).__init__(text, *args, **kw)


def get_box(img_arr):
    """Get the box of the digit. Return the left, top, right, bottom"""
    shape = img_arr.shape
    for i in range(shape[0]):
        if img_arr[i].max() == 255:
            top = i
            break

    for i in range(shape[0]-1, -1, -1):
        if img_arr[i].max() == 255:
            bottom = i
            break

    for i in range(shape[1]):
        if img_arr[:, i].max() == 255:
            left = i
            break

    for i in range(shape[1]-1, -1, -1):
        if img_arr[:, i].max() == 255:
            right = i
            break
    try:
        box = left, top, right, bottom
    except NameError as e:
        print("This shouldn't have happened.")
        raise ImageError

    return box


def pre_process(file_name):
    """Resize the image, move the digit to the center of image according its box.
    Finally return a PIL.Image object of 32x32 size."""
    img = Image.open(file_name)

    # When the picture comes from the front-end, only the Alpha band takes the information.
    # When the picture is generated from data set, this sentence take no use.
    img = img.convert('L')  # //
    # /////////////////////////

    imgs_arr, borders = split_digits(img)

    out = []

    for i in range(len(imgs_arr)):
        img_arr = imgs_arr[i]
        left, top, right, bottom = get_box(img_arr)
        # get_box() is designed for single digit pictures, at the early period.
        # But now we need to recognize pictures that has multiple digits, so the left and right should be rewritten.
        left = borders[i*2]
        right = borders[i*2+1]

        # stretch the image's height to standard IMG_SIZE
        size = (right - left, bottom - top)
        ratio = IMG_SIZE[1] / float(size[1])
        new_size = int(ratio * size[0]), IMG_SIZE[1]
        center = (new_size[0]/2, new_size[1]/2)
        rel = int(IMG_SIZE[0]/2 - center[0])  # x relative move of the center of IMG_SIZE and new_size

        print("center:%s ratio:%f size:%s new_size:%s" %
              (str(center), ratio, str(size), str(new_size)))
        im = img.crop((left, top, right, bottom))
        im = im.resize(new_size)

        img_temp = Image.new('L', IMG_SIZE, 0)
        img_temp.paste(im, (rel, 0))
        out.append(img_temp)

    return out


def get_array(img):
    """Convert the image to the numpy array."""
    arr = np.asarray(img).reshape([1, 1024])
    arr = arr/255

    return arr


def split_digits(img):
    """Split one picture with multiple digits into several pictures with individual digit.
    Return the numpy arrays
    """
    img = img.convert('L')
    img_arr = np.asarray(img, dtype=int)
    borders = []  # the x coordinate left and right of each digit

    # find the borders of each digit
    last_color = 0
    for x in range(len(img_arr[0])):
        this_color = img_arr[:, x].max()
        if last_color != this_color:
            borders.append(x)

        last_color = this_color

    rslt = []
    for i in range(0, len(borders), 2):
        left = borders[i]
        right = borders[i+1]
        rslt.append(img_arr[:, left:right])

    return rslt, borders


if __name__  == '__main__':
    np.set_printoptions(threshold=np.inf, linewidth=500)

    # im = Image.open(FILE_NAME)
    # im = im.resize((74, 42))
    # print(split_digits(im))

    import knn_recog

    ims = pre_process(FILE_NAME)
    for im in ims:
        arr = get_array(im)
        result = knn_recog.predict(arr)





